import numpy as np
import math
from timer import timer
from hist import hist

__dtype__ = 'float32'

class kmeans:
    data = None
    means = None
    indicators = None
    K = None
    count = None
    __dbg__ = None
    # hist = None
    
    def __init__(self, data_file, K, debug=False):
        ''' initialise new algorithm '''
        dat = np.load(data_file)
        self.data = dat['data']
        self.K = K
        self.__dbg__ = debug
        self.solve()
    
    def solve(self, runs=1000, tol=1e-5):
        ''' solve the k means algorithm '''
        self.means = np.random.choice(self.data, self.K).astype(__dtype__)
        if self.__dbg__:
            print 'starting solve'
        busy = True
        r = 1
        # self.hist = hist(self.data)
        while busy:
            if self.__dbg__:
                print 'run', r
            oldmeans = np.array(self.means)

            # neg_c, neg_b, pos_c, pos_b = self.hist
            tm = timer(0, 1)

            '''
            pos_diffs = np.square(np.abs([pos_b - self.means[i] for i in range(self.K)], dtype=__dtype__))
            if len(neg_c > 0):
                neg_diffs = np.square(np.abs([neg_b - self.means[i] for i in range(self.K)], dtype=__dtype__))

            tm.reset()
            '''
            diffs = np.square(np.abs([self.data - self.means[i] for i in range(self.K)], dtype=__dtype__))
            
            '''
            diffs = np.zeros((len(self.means), len(self.data)))
            sort = np.zeros((len(self.data)))
            count = 0
            for i in range(len(pos_c)):
                for j in range(len(self.means)):
                    diffs[j, count:count+pos_c[i]] = pos_diffs[j, i]
                sort[count:count+pos_c[i]] = i
                count += pos_c[i]

            for i in range(len(neg_c)):
                for j in range(len(self.means)):
                    diffs[j, count:count+neg_c[i]] = neg_diffs[j, i]
                sort[count:count+neg_c[i]] = -i - 1
                count += neg_c[i]
            '''

            self.indicators = np.argmin(diffs, axis=0)
            self.count = np.zeros_like(self.means)
            colors = ['r', 'b', 'bk']

            alldata = list()
            for i in range(self.K):
                cur = self.data[np.where(self.indicators == i)]
                self.means[i] = np.mean(cur)
                self.count[i] = len(cur)
                alldata.append(np.array(cur))

            if self.__dbg__:
                print self.means
                print self.count / np.sum(self.count)

            if len(self.count[np.where(self.count == 0)]) == 0:
                r += 1
                if r == runs:
                    busy = False
                
                if len(self.means[np.where(np.abs(self.means - oldmeans) > tol)]) == 0:
                    break
            else:
                if self.__dbg__:
                    print 'nan - new means'
                self.means = np.random.choice(self.data, self.K).astype(__dtype__)
                r = 0
    
        return 0

if __name__ == '__main__':
    from gmm import gmm
    # gmm_data = np.load('data.npz')
    # data = gmm_data['data'].astype('int')
    data = gmm.random_data(((1000, 0, 5),(1000, -10, 5),(100,25,1))).astype('int')
    np.savez('data.npz')
    km = kmeans('data.npz', 3, debug=True)
    # pp.show()

    print 
    print km.means
    print km.count
